import h5py
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import json
from tqdm import tqdm
import time
import os

def find_top_k_similar_batched(h5_path, top_k=20, batch_size=1024, output_json_path='similarity_results.json'):
    """
    通过分批处理，高效计算HDF5文件中每个嵌入向量的最相似项。
    适用于大规模数据集，避免内存溢出。

    :param h5_path: 输入的 HDF5 文件路径。
    :param top_k: 需要找出的最相似项的数量。
    :param batch_size: 每次处理的查询向量数量。可根据内存大小调整。
    :param output_json_path: 输出的 JSON 文件路径。
    """
    print("--- 开始大规模相似度计算 ---")
    
    # --- 1. 加载全体数据作为“数据库” ---
    print(f"正在从 '{h5_path}' 加载所有嵌入向量到内存...")
    try:
        with h5py.File(h5_path, 'r') as f:
            filenames = list(f.keys())
            num_embeddings = len(filenames)
            if num_embeddings == 0:
                print("错误：HDF5 文件为空。")
                return

            # 获取嵌入维度
            embedding_dim = f[filenames[0]].shape[0]
            
            # 预分配内存以提高效率
            embedding_matrix = np.zeros((num_embeddings, embedding_dim), dtype=np.float32)
            for i, key in enumerate(tqdm(filenames, desc="Loading embeddings")):
                embedding_matrix[i] = f[key][:]
    except Exception as e:
        print(f"加载HDF5文件时出错: {e}")
        return

    print(f"加载完成！共 {num_embeddings} 个向量，维度为 {embedding_dim}。")
    print(f"数据库矩阵内存占用: {embedding_matrix.nbytes / 1e9:.2f} GB")

    # --- 2. 分批计算与流式写入 ---
    # 估算单批次相似度矩阵的内存占用，给用户一个参考
    sim_batch_mem = (batch_size * num_embeddings * 4) / 1e9
    print(f"使用 batch_size = {batch_size}，每个批次的相似度矩阵内存占用约 {sim_batch_mem:.2f} GB。")
    if sim_batch_mem > 8: # 假设机器有16GB以上内存，给个警告
        print("警告：batch_size 可能过大，如果遇到内存问题，请尝试减小它。")

    # 使用流式写入，避免最终结果字典占用过多内存
    with open(output_json_path, 'w', encoding='utf-8') as f_out:
        f_out.write('{\n') # JSON 文件开始

        total_batches = (num_embeddings + batch_size - 1) // batch_size
        pbar = tqdm(range(0, num_embeddings, batch_size), total=total_batches, desc="Processing batches")
        
        is_first_entry = True

        for i in pbar:
            # 定义当前批次的范围
            start_index = i
            end_index = min(i + batch_size, num_embeddings)
            
            # 获取查询批次
            query_batch = embedding_matrix[start_index:end_index]
            query_filenames = filenames[start_index:end_index]

            # --- 核心计算 ---
            # 计算批次与整个数据库的相似度
            sim_batch = cosine_similarity(query_batch, embedding_matrix)
            
            # --- 高性能Top-K查找 ---
            # 使用 argpartition 找到每行 Top (k+1) 个值的索引 (包含自身)
            # - (top_k + 1) 表示我们想要最大的 k+1 个元素
            # [:, -(top_k+1):] 提取出这些元素的索引
            top_indices_batch = np.argpartition(sim_batch, -(top_k + 1), axis=1)[:, -(top_k + 1):]

            # --- 整理并写入结果 ---
            for j in range(len(query_filenames)):
                current_filename = query_filenames[j]
                
                # 获取这个向量的 Top-K 索引和分数
                top_indices = top_indices_batch[j]
                top_scores = sim_batch[j, top_indices]
                
                # 将索引和分数打包，并按分数降序排序 (只需排序 k+1 个元素，非常快)
                results_with_scores = sorted(zip(top_scores, top_indices), reverse=True)
                
                # 格式化并排除自身
                formatted_results = []
                for score, idx in results_with_scores:
                    if idx != (start_index + j): # 排除自身
                        formatted_results.append({filenames[idx]: float(score)})
                    if len(formatted_results) == top_k:
                        break
                
                # 写入到JSON文件
                if not is_first_entry:
                    f_out.write(',\n')
                
                json_entry = json.dumps({current_filename: formatted_results})
                # 去掉外层的 {}，因为我们是手动构建整个对象
                f_out.write(f'  "{current_filename}": {json.dumps(formatted_results)}')
                is_first_entry = False

        f_out.write('\n}\n') # JSON 文件结束

    print(f"--- 所有操作完成！结果已保存到 '{output_json_path}' ---")

# --- 执行主程序 ---
if __name__ == "__main__":
    H5_FILE = 'outputs/baseline_embedding/all_embeddings.h5' # 请替换为你的文件路径
    OUTPUT_FILE = 'similarity_results_batched.json'
    TOP_N = 20
    # **关键参数**：根据你的可用内存调整。
    # 如果内存有16GB，1024是个不错的起点。如果只有8GB，可以尝试512或256。
    BATCH_SIZE = 10000 

    # 检查输入文件是否存在
    if not os.path.exists(H5_FILE):
        print(f"错误: 输入文件 '{H5_FILE}' 不存在。请先运行上一个脚本创建示例文件或指定正确路径。")
    else:
        start_time = time.time()
        find_top_k_similar_batched(
            h5_path=H5_FILE, 
            top_k=TOP_N, 
            batch_size=BATCH_SIZE, 
            output_json_path=OUTPUT_FILE
        )
        end_time = time.time()
        print(f"总耗时: {end_time - start_time:.2f} 秒")
        
        
        
        
